'''PDBfileHelper has utility functions to get some information out of PDB files and also to modify it.'''

import os, logging, string, sys, re, shutil, math
sstrip = string.strip
from collections import OrderedDict

class PDBfileHelper :
	'''class for helping out with PDB file parsing etc'''
	def __init__(s, pdbfile) :
		s.pdbfile = pdbfile
		s.indices = {} # to store indices of fields in PDB lines
		s.indices['atomname'] = (12,16)
	def hasTLSgroups(s) :
		'''Returns True if there are no ANISOU records and TLS groups in PDB remarks are nonzero.'''
		rexp = re.compile("NUMBER *OF *TLS *GROUPS")
		rexp1 = re.compile(": *0")
		hasANISOU = False
		for l in open(s.pdbfile,'r') :
			if len(l) > 6 and l[0:6] == "ANISOU" : hasANISOU = True; break
		for l in open(s.pdbfile,'r') :
			if not rexp.search(l) : continue
			try :
				numtls = string.atoi(l.split(':')[1])
				if numtls <= 0 : return False
				if numtls > 0 and hasANISOU == True : return False
				if numtls > 0 and hasANISOU == False : return True
			except : return False
		return False
	def getResolution(s) :
		'''Reads resolution from REMARK 2 RESOLUTION line in pdbfile, else returns None.'''
		resolStr = "REMARK   2 RESOLUTION"
		resolStrlen = len(resolStr)
		for l in open(s.pdbfile,'r') :
			if len(l) > resolStrlen and l[0:resolStrlen] == resolStr :
				try : return string.atof(l[23:30])
				except : return None
		return None
	def getRfactors(s) :
		'''parses R factors from remark 3, returns None for missing R factor'''
		if not hasattr(s,'r') :
			Rre = re.compile("REMARK *3 *R *VALUE *.WORKING *SET. *:") ## specific
			RFREEre = re.compile("REMARK *3 *FREE *R *VALUE *:") ## specific
			r, rf = None, None
			for l in open(s.pdbfile,'r') :
				if Rre.search(l) :
					try : r = string.atof(l.split(":")[1])
					except : pass
				if RFREEre.search(l) :
					try : rf = string.atof(l.split(":")[1])
					except : pass
			if r!= None and rf != None : s.r, s.rf = r, rf
			Rre = re.compile("REMARK *3 *R *VALUE.*:") ## any kind of R 
			RFREEre = re.compile("REMARK *3 *FREE *R *VALUE.*:") ## any kind of Rfree
			r1, rf1 = None, None
			for l in open(s.pdbfile,'r') :
				if not r1 and Rre.search(l) :
					try : r1 = string.atof(l.split(":")[1])
					except : pass
				if not rf1 and RFREEre.search(l) :
					try : rf1 = string.atof(l.split(":")[1])
					except : pass
			if r != r1 and r == None : r = r1
			if rf != rf1 and rf == None : rf = rf1
			s.r, s.rf = r, rf
		return s.r, s.rf
	def getSpacegroup(s) :
		if not hasattr(s,'spacegroup') : s.parseCryst1()
		return s.spacegroup
	def getCellStr(s) :
		'''string containing a b c alpha beta gamma'''
		return string.join([sstrip("%6.3f"%cellelem) for cellelem in s.getCell()], " ")
	def getCell(s) :
		if not hasattr(s,'cell') : s.parseCryst1()
		return s.cell
	def parseCryst1(s) :
		'''Returns a list of 6 numbers from CRYST1 card'''
		s.spacegroup, s.cell = None, None
		for l in open(s.pdbfile,'r') :
			if len(l) <= 6 or l[0:6] != "CRYST1" : continue
			s.spacegroup = re.sub("\s+$", "", re.sub("^\s+", "", l[55:66]))
			s.cell = ( string.atof(l[ 6 : 15 ]), string.atof(l[ 15 : 24 ]), string.atof(l[ 24 : 33 ]),
			 	string.atof(l[ 33 : 40 ]), string.atof(l[ 40 : 47 ]), string.atof(l[ 47 : 54 ]), )
			break
	def countNCSoperators(s) :
		'''count num of MTRIX operators in a pdb file'''
		numncs = 0
		for l in open(s.pdbfile,'r') :#.readlines() :
			if "MTRIX" in l[0:5] : numncs += 1
		if numncs % 3 != 0 :
			logging.error("Number of MTRIX lines in " + s.pdbfile + " = " + numncs + " : not multiple of 3! Aborting....") ; sys.exit(0)
		return numncs / 3
	def getDepositionDate(s) :
		'''from HEADER, return full date and year. else None, None'''
		for l in open(s.pdbfile,'r') :#.readlines() :
			if "HEADER" in l[0:6] and len(l) > 57 :
				depdate = l[50:59] ; depyear = depdate[7:]
				return depdate, depyear
		return None, None
	def countNumAtoms(s) :
		'''return sum total of ATOM and HETATM records, #ATOMs, #HETATMs, #Calpha atoms, #UNK/X'''
		numatom, numhet, numCA, numUNK = 0, 0, 0, 0
		for l in open(s.pdbfile,'r') :#.readlines() :
			if "ATOM" in l[0:4] :
				numatom += 1
				if ' CA ' in l[12:16] : numCA += 1
				if 'UNK UN' in l or 'UNX UN' in l : numUNK += 1 
			if "HETATM" in l[0:6] :
				numhet += 1
				if 'UNK UN' in l or 'UNX UN' in l : numUNK += 1 
		return numatom+numhet, numatom, numhet, numCA, numUNK
	def getRefinementProgram(s) :
		'''from REMARK-3. else None'''
		for l in open(s.pdbfile,'r') :#.readlines() :
			if 'REMARK   3   PROGRAM' in l[0:20] :
				l = re.sub("\n","",l)
				l = re.sub('REMARK *3 *PROGRAM *:', '', l)
				l = re.sub( '\s+$', '', re.sub('^\s+','',l) )
				return l
		return None
	def correctANISOatomnum(s) :
		'''make atom numbers on ATOM/HETATM and subsequent ANISOU line identical if they are not. (modifies the pdbfile!)'''
		lines = open(s.pdbfile,'r').readlines()
		ofp = open(s.pdbfile, 'w')
		for li in range(len(lines)-1) :
			if lines[li+1][0:6] == "ANISOU" and lines[li][0:6] in ["ATOM  ","HETATM"] and lines[li][6:11] != lines[li+1][6:11] :
				ofp.write(lines[li][0:6] + lines[li+1][6:11] + lines[li][11:])
			else : ofp.write(lines[li])
		ofp.close()
	def correctMTRIXrepeats(s) :
		'''if MTRIXi records are repeated in consecutive lines, remove the latter line. (modifies the pdbfile!)'''
		lines = open(s.pdbfile,'r').readlines()
		ofp = open(s.pdbfile, 'w')
		for li in range(len(lines)) :
			if lines[li][0:5] == "MTRIX" and lines[li-1][0:6] == lines[li][0:6] : pass
			else : ofp.write(lines[li])
		ofp.close()
	def removeUNKatoms(s) :
		'''remove UNK UNX UNL atoms irrespective of residues they occur in'''
		lines = open(s.pdbfile,'r').readlines()
		ans,ane = s.indices['atomname']
		unkatomnames = ["UNK ","UNX "," UNX"," UNK"]
		ofp = open(s.pdbfile, 'w')
		for li in range(len(lines)) :
			if lines[li][0:6] in ["ATOM  ","HETATM"] and lines[li][ans:ane] in unkatomnames : continue
			ofp.write(lines[li])
		ofp.close()

